"港航 AI 识别完毕后处理代码"

__author__ = "zwj"

# 导入根目录作为模块搜索路径
import os, sys

# ([".."] * 4) 就是项目目录里面 4 层
parent_path = os.path.abspath(os.path.join(__file__, *([".."] * 4)))
sys.path.insert(0, parent_path)

import yaml, time, cv2
from typing import Dict, List
from enum import Enum, unique
import math
from collections import deque

from core.task.modules.processors.base_processor import BaseProcessor
from core.utils.state_machine import StateMachine
from core.utils.logger import logging

from projects.GangHang.processors.http_message import HttpMessage, CameraType, FocusMode


# 船只航向枚举
@unique
class Heading(Enum):
    kLeft = "left"
    kRight = "right"
    kUnknown = "unknown"


# 像素点坐标
class PointData:
    def __init__(self, x=0, y=0):
        # 确保点坐标都为正数
        self.x = abs(x)
        self.y = abs(y)


# 船只检测数据汇总
class ShipData:
    def __init__(
        self,
        cls_id,
        track_id,
        ship_start: PointData,
        ship_end: PointData,
    ):
        # 船只检测
        self.cls_id = cls_id
        self.track_id = track_id
        self.ship_start = ship_start
        self.ship_end = ship_end

    def get_midpoint_data(self) -> PointData:
        return PointData(
            (self.ship_start.x + self.ship_end.x) / 2,
            (self.ship_start.y + self.ship_end.y) / 2,
        )

    def get_area(self) -> int:
        return (self.ship_end.x - self.ship_start.x) * (
            self.ship_end.y - self.ship_start.y
        )


# 船只检测数据
class Mask:
    def __init__(self, width, height):
        self.__width = width
        self.__height = height
        # 初始化一个二维列表
        self.clear()

    def set_mask(self, start_point: PointData, end_point: PointData):
        # 确保数据不过界
        start_point = self.__point_data_int(start_point)
        end_point = self.__point_data_int(end_point)

        # 将给定框范围内的元素置为True
        for y in range(start_point.y, end_point.y):
            for x in range(start_point.x, end_point.x):
                self.__mask[y][x] = True

    def get_mask(self, point: PointData):
        point = self.__point_data_int(point)
        return self.__mask[int(point.y)][int(point.x)]

    def clear(self):
        self.__mask = [
            [False for _ in range(self.__width)] for _ in range(self.__height)
        ]

    def __point_data_int(self, point: PointData):
        point_x = min(int(point.x), self.__width - 1)
        point_y = min(int(point.y), self.__height - 1)
        return PointData(int(point_x), int(point_y))


@unique
class StateMachineEnum(Enum):
    kDetect = 0  # 检测船只,锁定要处理的船只
    kZoom = 1  # 变焦
    kZoomDetect = 2  # 变焦后检测


class GangHangProcessor(BaseProcessor):
    def __init__(
        self,
        keys,
        Custom_cls_map,
        Screen_size,
        Left_Heading,
        HttpMessage_cfg_path,
    ) -> None:

        # * key 值对应索引
        self.__keys = keys

        # * 输入数据
        self.__in_key = self.__keys["in"]
        self.__track_key = self.__keys["track"]
        self.__out_key = self.__keys["out"]
        self.__finish_tag_key = self.__keys["finish_tag"]
        self.__mqtt_key = self.__keys["mqtt"]

        # * 检测 id 对应含义地图
        if isinstance(Custom_cls_map, str):
            with open(Custom_cls_map, "r") as f:
                self.__cls_map = {
                    x: y for x, y in enumerate(yaml.safe_load(f)["label_list"])
                }
        else:
            self.__cls_map = Custom_cls_map

        # * 绘图框颜色初始化, BGR 颜色, 0 0 0 为黑色
        self.__color = (0, 0, 255)

        # * 图片输出目录初始化
        self.__image_output_path = os.getcwd() + "/output"
        if not os.path.exists(self.__image_output_path):
            os.makedirs(self.__image_output_path)

        # * 获取画面大小
        self.__screen_size_x = Screen_size[0]
        self.__screen_size_y = Screen_size[1]

        # * 状态机
        self.__state_machine = StateMachine()
        self.__state_machine.set_state_config(StateMachineEnum.kDetect, self.__detect)
        self.__state_machine.set_state_config(StateMachineEnum.kZoom, self.__zoom)
        self.__state_machine.set_state_config(
            StateMachineEnum.kZoomDetect, self.__zoom_detect
        )
        self.__state_machine.set_initial_state(StateMachineEnum.kDetect)

        # * 船只数据
        self.__ship_data_list: List[ShipData] = []
        self.__detecting_ships_list: List[ShipData] = []

        self.__ship_offset_map = {}
        self.__ship_offset_threshold = 4

        self.__track_id = 0
        self.__track_id_midpoint = PointData(0, 0)

        self.__zoom_sleep_count = 0
        self.__zoom_sleep_threshold = 10
        self.__zoom_x_deviation_threshold = 0.2

        # * 船只变焦下截图
        self.__zoom_image_flag = False
        self.__zoom_image_time = time.time()

        # 船只掩码
        self.__mask = Mask(self.__screen_size_x, self.__screen_size_y)

        # 船只航向要求
        self.__set_ship_heading(Left_Heading)

        # * 控制接口初始化
        self.__http_message = HttpMessage()
        self.__http_message.init_cfg(HttpMessage_cfg_path)

        # * DEBUG : 测试整个追踪结果
        self.__debug = False

        # * 记录操作时间戳
        self.__last_fps_time = time.time()
        self.__last_mqtt_time = time.time()

    def __set_ship_heading(self, Left_Heading):
        # 初始化一个掩码二维列表,将不检测区屏蔽
        self.__mask.clear()
        # self.__mask.set_mask(
        #     PointData(0, 0),
        #     PointData(self.__screen_size_x, self.__screen_size_y // 3),
        # )

        if Left_Heading:
            self.__ship_heading = Heading.kLeft
            self.__cls_id = 2

            self.__mask.set_mask(
                PointData(self.__screen_size_x // 2, 0),
                PointData(self.__screen_size_x, self.__screen_size_y),
            )
        else:
            self.__ship_heading = Heading.kRight
            self.__cls_id = 0

            self.__mask.set_mask(
                PointData(0, 0),
                PointData(self.__screen_size_x // 2, self.__screen_size_y),
            )

    def __refresh_ship_data(self, data: Dict) -> Dict:
        # 获取结果
        track_result = data[self.__track_key]

        ship_data_list = []
        detecting_ships_list = []

        for bboxes, track_ids in zip(track_result["boxes"], track_result["track_ids"]):

            for box, track_id in zip(bboxes, track_ids):
                cls_id, score, x1, y1, x2, y2 = box
                cls_id = int(cls_id)

                # 只记录检测区域内船只
                ship_data = ShipData(
                    cls_id, track_id, PointData(x1, y1), PointData(x2, y2)
                )

                ship_data_list.append(ship_data)

                if cls_id != self.__cls_id:
                    continue

                # 获取到的船中心点如果在掩码区域内跳过
                if self.__mask.get_mask(ship_data.get_midpoint_data()):
                    continue

                detecting_ships_list.append(ship_data)

        self.__ship_data_list = ship_data_list
        self.__detecting_ships_list = detecting_ships_list

        return data

    def __detect_image(self, data: Dict, ships_list: List[ShipData]) -> Dict:
        image = data[self.__in_key][0].copy()
        draw_thickness = min(image.shape[:2]) // 320

        if self.__zoom_image_flag and time.time() - self.__zoom_image_time > 1:
            self.__zoom_image_time = time.time()
            cv2.imwrite(f"{self.__image_output_path}/{time.time()}.jpg", image)

        mqtt_message = {}
        mqtt_message["id"] = data["id"]
        mqtt_ship_data = []

        for ship_data in ships_list:
            # 画出追踪 ID
            cv2.putText(
                image,
                str(ship_data.track_id),
                (int(ship_data.ship_start.x), int(ship_data.ship_start.y)),
                cv2.FONT_HERSHEY_COMPLEX,
                1,
                self.__color,
                2,
            )
            # 画出检测类型
            cv2.putText(
                image,
                str(ship_data.cls_id),
                (int(ship_data.ship_start.x), int(ship_data.ship_end.y)),
                cv2.FONT_HERSHEY_COMPLEX,
                1,
                (0, 0, 0),
                2,
            )
            # 画出追踪框
            image = cv2.rectangle(
                image,
                (int(ship_data.ship_start.x), int(ship_data.ship_start.y)),
                (int(ship_data.ship_end.x), int(ship_data.ship_end.y)),
                color=self.__color,
                thickness=draw_thickness,
            )

            ship_json = {
                "track_id": ship_data.track_id,
                "cls_id": ship_data.cls_id,
                "x1": ship_data.ship_start.x / self.__screen_size_x,
                "y1": ship_data.ship_start.y / self.__screen_size_y,
                "x2": ship_data.ship_end.x / self.__screen_size_x,
                "y2": ship_data.ship_end.y / self.__screen_size_y,
            }

            if self.__track_id != 0 and self.__track_id == ship_data.track_id:
                image = cv2.circle(
                    image,
                    (int(self.__track_id_midpoint.x), int(self.__track_id_midpoint.y)),
                    radius=10,
                    color=self.__color,
                    thickness=-1,
                )
                ship_json["target"] = True
            else:
                ship_json["target"] = False

            mqtt_ship_data.append(ship_json)

        mqtt_message["data"] = mqtt_ship_data

        # 限制发送帧率 3 FPS
        if time.time() - self.__last_mqtt_time > 1 / 3:
            self.__last_mqtt_time = time.time()

            # 推送追踪船只信息
            data[self.__mqtt_key] = mqtt_message

        data[self.__out_key] = image
        return data

    def __detect(self):
        # 根据航向对船只进行排序:
        if self.__ship_heading == Heading.kLeft:
            # 从左到右
            self.__detecting_ships_list = sorted(
                self.__detecting_ships_list, key=lambda box: (box.ship_start.x)
            )
        elif self.__ship_heading == Heading.kRight:
            # 从右到左
            self.__detecting_ships_list = sorted(
                self.__detecting_ships_list, key=lambda box: (-box.ship_end.x)
            )

        # 追踪的船只消失了
        if self.__track_id != 0 and self.__track_id not in [
            x.track_id for x in self.__detecting_ships_list
        ]:
            logging.info(f"锁定船只 {self.__track_id} 失踪")
            self.__track_id = 0

        for ship_data in self.__detecting_ships_list:
            midpoint_data = ship_data.get_midpoint_data()
            y = midpoint_data.y
            y_deviation = y / self.__screen_size_y - 0.5
            adjust_the_angle = 100 * y_deviation

            if not ship_data.track_id in self.__ship_offset_map:
                self.__ship_offset_map[ship_data.track_id] = deque(maxlen=5)
            ship_processing = self.__ship_offset_map[ship_data.track_id]

            ship_processing.append(adjust_the_angle)

            # 跳过非锁定船只
            if self.__track_id != 0 and self.__track_id != ship_data.track_id:
                continue

            if len(ship_processing) == ship_processing.maxlen:
                # 第一个满足条件的追踪
                self.__track_id = ship_data.track_id

                adjust_the_angle = sum(ship_processing) / len(ship_processing)

                logging.info(
                    f"锁定船只 {ship_data.track_id} : {y_deviation} {adjust_the_angle}"
                )
                break

        if self.__track_id == 0:
            return StateMachineEnum.kDetect, 0

        self.__ship_offset_map.clear()

        self.__track_id_midpoint = midpoint_data

        # 误差调整,让图像整体往下走
        # adjust_the_angle += 5

        if abs(adjust_the_angle) < self.__ship_offset_threshold:
            # 计算船头偏差
            # x_deviation = midpoint_data.x / self.__screen_size_x - 0.5

            if self.__ship_heading == Heading.kLeft:
                x_deviation = ship_data.ship_end.x / self.__screen_size_x - 0.5
            elif self.__ship_heading == Heading.kRight:
                x_deviation = ship_data.ship_start.x / self.__screen_size_x - 0.5

            if abs(x_deviation) > 0.12:
                logging.info(
                    f"船只 {ship_data.track_id}  偏离中心位置 x : {x_deviation}, y : {adjust_the_angle}, 等待进入..."
                )
                return StateMachineEnum.kDetect, 0

            target_width = self.__screen_size_x / (
                ship_data.ship_end.x - ship_data.ship_start.x
            )
            target_height = self.__screen_size_y / (
                ship_data.ship_end.y - ship_data.ship_start.y
            )
            # 计算变焦倍率
            self.__zoom_factor = max(target_width, target_height)
            # 设置变焦最小值 ,避免数字变焦
            self.__zoom_factor = int(max(self.__zoom_factor, 7))

            logging.info(
                f"船只 {ship_data.track_id} 处于中心位置 {adjust_the_angle}, 变焦倍率 : {self.__zoom_factor},进入变焦检测逻辑..."
            )
            self.__track_id = 0
            return StateMachineEnum.kZoom, 1

        logging.info(f"云台转动 : {adjust_the_angle}")
        self.__http_message.camera_screen_drag(-adjust_the_angle)
        return StateMachineEnum.kDetect, 2

    def __zoom(self):
        logging.info("切换到变焦镜头,镜头缩放...")

        # 循环调用接口，直到成功或者失败5次
        for _ in range(5):
            if self.__http_message.camera_focal_length_set(
                CameraType.kZoom, self.__zoom_factor
            ):
                break

        for _ in range(5):
            if self.__http_message.live_lens_change(CameraType.kZoom):
                break

        # self.__http_message.camera_focal_length_set(
        #     CameraType.kZoom, self.__zoom_factor
        # )
        # self.__http_message.live_lens_change(CameraType.kZoom)
        self.__zoom_sleep_count = 0
        return StateMachineEnum.kZoomDetect, 2

    def __zoom_detect(self):
        adjust_the_angle = 0
        self.__zoom_image_flag = True

        # 等待船只入画
        if len(self.__ship_data_list) == 0:
            # logging.info(f"变焦未检测到船只,进行等待...")
            self.__zoom_sleep_count += 1
            logging.info(f"变焦下船只 {self.__track_id} 丢失,加速退出变焦")
        # 中间有检测到一次就继续保持数秒
        else:
            # 找到最大的框
            max_area = 0
            max_ship_data = None
            track_id = 0

            # 极小概率 BUG : 前一个 id 消失,然后有一个最大的框出现?
            for ship_data in self.__ship_data_list:
                # 有追踪目标提前退出
                if self.__track_id == ship_data.track_id:
                    max_ship_data = ship_data
                    break

                area = ship_data.get_area()
                if area > max_area:
                    max_area = area
                    max_ship_data = ship_data
                    track_id = ship_data.track_id

            if self.__track_id == 0:
                self.__track_id = track_id
                logging.info(f"变焦下开始追踪船只 {self.__track_id}...")

            self.__track_id_midpoint = max_ship_data.get_midpoint_data()
            # 计算船中心点偏移量
            x_deviation = self.__track_id_midpoint.x / self.__screen_size_x - 0.5
            y_deviation = self.__track_id_midpoint.y / self.__screen_size_y - 0.5
            adjust_the_angle = 100 * y_deviation

            if self.__ship_heading == Heading.kLeft:
                # 0.4 > 0.3
                x_deviation_flag = x_deviation > self.__zoom_x_deviation_threshold
            elif self.__ship_heading == Heading.kRight:
                # -0.4 < -0.3
                x_deviation_flag = x_deviation < -self.__zoom_x_deviation_threshold

            logging.info(
                f"变焦下船只 {self.__track_id} x : {x_deviation} , y : {adjust_the_angle}"
            )

            if abs(y_deviation) > 0.3 or x_deviation_flag:
                self.__zoom_sleep_count += 2
                logging.info(
                    f"变焦下船只 {self.__track_id} {x_deviation} {y_deviation} 偏移过大加速退出变焦"
                )
            else:
                self.__zoom_sleep_count = 0

        if self.__zoom_sleep_count < self.__zoom_sleep_threshold:
            # 变焦镜头下移动下镜头,需要增大阈值防止来回调整
            if abs(adjust_the_angle) > self.__ship_offset_threshold * 2:
                logging.info(f"变焦下云台转动 : {adjust_the_angle}")
                self.__http_message.camera_screen_drag(-adjust_the_angle)
                self.__zoom_sleep_count = 0
                return StateMachineEnum.kZoomDetect, 1

            return StateMachineEnum.kZoomDetect, 0

        # 等待镜头还原
        logging.info("恢复广角镜头...")
        self.__http_message.live_lens_change(CameraType.kWide)
        self.__http_message.get_recovery_camera_pitch()
        self.__track_id = 0
        self.__zoom_image_flag = False
        return StateMachineEnum.kDetect, 1

    def _call(self, data: Dict) -> Dict:
        # 读取完毕关闭任务
        if data.get(self.__finish_tag_key, False):
            logging.info("任务完成")
            self.close()

        # 计算出调用该函数的帧率
        # fps = 1 / (time.time() - self.__last_fps_time)
        # logging.info(f"帧率 : {fps}")
        # self.__last_fps_time = time.time()

        # 获取船只数据
        data = self.__refresh_ship_data(data)

        # 运行状态机
        self.__state_machine.run_once()

        # 画出状态机流程的检测框
        # data = self.__detect_image(data, self.__detecting_ships_list)

        # 画出所有检测船只
        data = self.__detect_image(data, self.__ship_data_list)

        return data

    def _init_check(self):
        """
        任务执行前初始化检查参数
        """

        if self.__debug:
            mkv_path = os.path.join(self.__image_output_path, "output.mkv")
            self.__out = cv2.VideoWriter(
                mkv_path,
                cv2.VideoWriter_fourcc(*"XVID"),
                30,
                (self.__screen_size_x, self.__screen_size_y),
            )

        logging.info(f"设置变焦镜头对焦模式为 : AFC")
        self.__http_message.camera_focus_mode_set(CameraType.kZoom, FocusMode.kAFC)

    def _set(self, data: Dict):
        if data.get("Left_Heading", False):
            self.__set_ship_heading(data["Left_Heading"])

    def _start(self):
        pass
        # self._init_check()

    def _close(self):
        if self.__debug:
            self.__out.release()


if __name__ == "__main__":
    pass
